package backend.opt;

import backend.Instruction.RISCVbinary;
import backend.Instruction.RISCVload;
import backend.Instruction.RISCVmove;
import backend.Instruction.RISCVstore;
import backend.component.RISCVblock;
import backend.component.RISCVfunction;
import backend.Instruction.RISCVinst;
import backend.component.RISCVimm;
import backend.component.RISCVmodule;
import backend.Register.*;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

import static backend.Register.initReg.SP;
import static backend.Register.initReg.phyRegs;

public class CalleesaveRegister {
    private RISCVmodule module;
    private int changesize = 0;

    public CalleesaveRegister(RISCVmodule module) {
        this.module = module;
    }

    public void run(){
        dealwithcall();
    }

    public void dealwithcall(){
        for(RISCVfunction func : module.getFunctions()){
            if(!"main".equals(func.getName())) {
                //不是main函数，说明是调用的函数
                HashMap<RISCVreg, Integer> sreg;
                sreg = new HashMap<>();
                for (RISCVblock block : func.getBlocks()) {
                    for (RISCVinst inst : block.getInstructionlist()) {
                        if (isSRegister(inst.getReg())) {
                            //尝试给其分配空间存储
                            sreg.put(inst.getReg(), func.getStacksize());
                            //System.out.print(inst.getReg());
                            //System.out.println("   " + sreg.get(inst.getReg()));
                            func.addAllocSize(8);
                            changesize += 8;
                        }
                    }
                }
                int z = 0;
                RISCVblock block = func.getFirstBlock();
                for (RISCVinst inst : block.getInstructionlist()) {
                    if (inst instanceof RISCVbinary) {
                        RISCVbinary bin = (RISCVbinary) inst;
                        if (SP.equals(bin.getDst()) && SP.equals(bin.getSrc1())) {
                            if (!func.isBig()) {
                                //此时不存在 move 指令
                                int off = func.getStacksize() + changesize;
                                RISCVoperand offset;
                                if (off >= -2048 && off <= 2047) {
                                    offset = new RISCVimm(off, true);
                                    RISCVinst IIII = block.getInstructionlist().getFirst();
                                    ((RISCVbinary) IIII).setSrcRight(offset);
                                    z = 1;
                                } else {
                                    RISCVmove move = new RISCVmove(Arrays.asList(phyRegs.get(5), new RISCVimm(off, false)));
                                    offset = move.getDst();
                                    RISCVinst IIII = block.getInstructionlist().getFirst();
                                    ((RISCVbinary) IIII).setSrcRight(offset);
                                    block.addInstructionlistatHead(move);
                                    z = 2;
                                }
                            }
                            else{
                                //存在 move指令
                                int off = func.getAllocSize() + func.getArgsSize() + changesize;
                                RISCVinst vinst = block.findInstrbefore(bin);
                                if(vinst instanceof RISCVmove){
                                    ((RISCVmove)vinst).setSrc(new RISCVimm(off, false));
                                }
                                z = 2;
                            }
                            break;
                        }
                    }
                }
                addSregSave(sreg, block, z);

                block = func.getExitBlock();
                int siz = block.getInstructionlist().size();
                for (RISCVinst inst : block.getInstructionlist()) {
                    if (inst instanceof RISCVbinary) {
                        RISCVbinary bin = (RISCVbinary) inst;
                        if (SP.equals(bin.getDst()) && SP.equals(bin.getSrc1())) {
                            if (!func.isBig()) {
                                //此时不存在 move 指令
                                z = siz;
                                int off = func.getStacksize() + changesize;
                                RISCVoperand offset;
                                if (off >= -2048 && off <= 2047) {
                                    offset = new RISCVimm(off, true);
                                    RISCVinst IIII = block.getInstructionlist().get(z - 2);
                                    ((RISCVbinary) IIII).setSrcRight(offset);
                                    z = z - 2;
                                } else {
                                    RISCVmove move = new RISCVmove(Arrays.asList(phyRegs.get(5), new RISCVimm(off, false)));
                                    offset = move.getDst();
                                    RISCVinst IIII = block.getInstructionlist().get(z - 2);
                                    ((RISCVbinary) IIII).setSrcRight(offset);
                                    block.addInstructionlistatindex(move,z - 2);
                                    z = z - 2;
                                }
                            }
                            else{
                                //存在 move指令
                                int off = func.getAllocSize() + func.getArgsSize() + changesize;
                                RISCVinst vinst = block.getInstructionlist().get(z - 2);
                                if(vinst instanceof RISCVmove){
                                    ((RISCVmove)vinst).setSrc(new RISCVimm(off, false));
                                }
                                z = z - 2;
                                addSregLoad(sreg, block, z);
                            }
                            break;
                        }
                    }
                }
                addSregLoad(sreg, block, z);

            }
        }
    }
    private void addSregSave( HashMap<RISCVreg,Integer> sreg,RISCVblock block, int z) {
        for (Map.Entry<RISCVreg, Integer> entry : sreg.entrySet()) {
            // 获取键和值
            RISCVreg key = entry.getKey();
            Integer value = entry.getValue();
            String typ = key instanceof physicalReg?"sd":"fsw";
            if (value >= -2048 && value <= 2047) {
                RISCVstore store = new RISCVstore(Arrays.asList(key, initReg.SP, new RISCVimm(value, true)), typ);
                block.addInstructionlistatindex(store,z);
            } else {
                //处理大数情况
                RISCVmove move = new RISCVmove(Arrays.asList(phyRegs.get(5),new RISCVimm(value, false)));
                RISCVbinary add = new RISCVbinary("add", Arrays.asList(phyRegs.get(5), initReg.SP, move.getDst()));
                RISCVstore store = new RISCVstore(Arrays.asList(key, add.getDst(), new RISCVimm(0, true)), typ);
                block.addInstructionlistatindex(store,z);
                block.addInstructionlistatindex(add,z);
                block.addInstructionlistatindex(move,z);
            }

        }
    }


    private void addSregLoad( HashMap<RISCVreg,Integer> sreg,RISCVblock block, int z) {
        for (Map.Entry<RISCVreg, Integer> entry : sreg.entrySet()) {
            // 获取键和值
            RISCVreg key = entry.getKey();
            Integer value = entry.getValue();
            String typ = key instanceof physicalReg?"ld":"flw";
            if (value >= -2048 && value <= 2047) {
                RISCVload store = new RISCVload(Arrays.asList(key, initReg.SP, new RISCVimm(value, true)), typ);
                block.addInstructionlistatindex(store,z);
            } else {
                //处理大数情况
                RISCVmove move = new RISCVmove(Arrays.asList(phyRegs.get(5),new RISCVimm(value, false)));
                RISCVbinary add = new RISCVbinary("add", Arrays.asList(phyRegs.get(5), initReg.SP, move.getDst()));
                RISCVload store = new RISCVload(Arrays.asList(key, add.getDst(), new RISCVimm(0, true)), typ);
                block.addInstructionlistatindex(store,z);
                block.addInstructionlistatindex(add,z);
                block.addInstructionlistatindex(move,z);
            }

        }
    }

    private boolean isSRegister(RISCVreg reg){
        if(reg instanceof physicalReg){
            return ((physicalReg) reg).isSReg();
        }
        else if(reg instanceof physicalFReg){
            return ((physicalFReg) reg).isSReg();
        }
        return false;
    }
}
